/* Copyright (c) 2024 Huawei Technologies Co., Ltd.
* This file is a part of the CANN Open Software.
* Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
* Please refer to the License for details. You may not use this file except in compliance with the License.
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
* See LICENSE in the root of the software repository for the full text of the License.
* ===================================================================================================================*/
#include <gtest/gtest.h>
#include "ascend_ir/ascend_reg_ops.h"
#include "ascir_ops.h"
#include "ascend_ir/ascend_ir.h"
#include "ascend_ir/expression.h"
#include <iostream>
#include "graph/utils/node_utils_ex.h"

class UtestAscendIR : public testing::Test {
 protected:
  void SetUp() {}

  void TearDown() {}
};
using namespace ge;
using ge::expression::ExpressionPtr;
using ge::expression::Expression;

TEST_F(UtestAscendIR, TilingKey_OK) {
  AscGraph graph("test_graph");
  graph.SetTilingKey(10);
  EXPECT_EQ(graph.GetTilingKey(), 10);
}

TEST_F(UtestAscendIR, CreateAxis) {
  AscGraph graph("test_graph");
  const ExpressionPtr s0 = Expression::CreateExpression("s0");
  Axis &s0_axis = graph.CreateAxis("S0", s0);
  s0_axis.align = 10;
  auto s0_axis_find = graph.FindAxis(s0_axis.id);
  EXPECT_NE(s0_axis_find, nullptr);
  EXPECT_EQ(s0_axis_find->name, "S0");
  EXPECT_EQ(s0_axis_find->align, 10);

  auto axis_invalid = graph.FindAxis(-1);
  EXPECT_EQ(axis_invalid, nullptr);
}

TEST_F(UtestAscendIR, BlockSplit) {
  AscGraph graph("test_graph");
  const ExpressionPtr s0 = Expression::CreateExpression("s0");
  Axis &s0_axis = graph.CreateAxis("S0", s0);
  auto split_axis = graph.BlockSplit(s0_axis.id);
  EXPECT_NE(split_axis.first, nullptr);
  EXPECT_NE(split_axis.second, nullptr);
  auto &outer_axis = *split_axis.first;
  auto &inner_axis = *split_axis.second;
  EXPECT_EQ(inner_axis.type, Axis::AXIS_TYPE_BLOCK_INNER);
  EXPECT_EQ(outer_axis.type, Axis::AXIS_TYPE_BLOCK_OUTER);
}

TEST_F(UtestAscendIR, TileSplit) {
  AscGraph graph("test_graph");
  const ExpressionPtr s0 = Expression::CreateExpression("s0");
  Axis &s0_axis = graph.CreateAxis("S0", s0);
  s0_axis.align = 10;
  auto split_axis = graph.TileSplit(s0_axis.id);
  EXPECT_NE(split_axis.first, nullptr);
  EXPECT_NE(split_axis.second, nullptr);
  auto &outer_axis = *split_axis.first;
  auto &inner_axis = *split_axis.second;
  EXPECT_EQ(inner_axis.type, Axis::AXIS_TYPE_TILE_INNER);
  EXPECT_EQ(outer_axis.type, Axis::AXIS_TYPE_TILE_OUTER);
}

TEST_F(UtestAscendIR, MergeAxis) {
  AscGraph graph("test_graph");
  const ExpressionPtr s0 = Expression::CreateExpression("s0");
  Axis &s0_axis = graph.CreateAxis("S0", s0);
  const ExpressionPtr s1 = Expression::CreateExpression("s1");
  Axis &s1_axis = graph.CreateAxis("S1", s1);
  auto merge_axis = graph.MergeAxis({s0_axis.id, s1_axis.id});
  EXPECT_NE(merge_axis, nullptr);
  EXPECT_EQ(merge_axis->type, Axis::AXIS_TYPE_MERGED);
}


TEST_F(UtestAscendIR, ApplySplit) {
  auto ONE = Expression::CreateExpression(1);
  AscGraph graph("test_graph");
  const ExpressionPtr s0 = Expression::CreateExpression("s0");
  Axis &s0_axis = graph.CreateAxis("S0", s0);
  ascir_op::Data data("data", graph);
  auto data_node = ge::NodeUtilsEx::GetNodeFromOperator(data);
  EXPECT_NE(data_node, nullptr);
  data.attr.sched.exec_order = 1;
  data.attr.sched.axis = {s0_axis.id};
  data.y.dtype = ge::DT_FLOAT16;
  data.y.format = ge::FORMAT_ND;
  data.y.axis = {s0_axis.id};
  data.y.repeats = {s0};
  data.y.strides = {ONE};

  auto split_axis = graph.TileSplit(s0_axis.id);
  EXPECT_NE(split_axis.first, nullptr);
  EXPECT_NE(split_axis.second, nullptr);
  auto &outer_axis = *split_axis.first;
  auto &inner_axis = *split_axis.second;
  EXPECT_EQ(inner_axis.type, Axis::AXIS_TYPE_TILE_INNER);
  EXPECT_EQ(outer_axis.type, Axis::AXIS_TYPE_TILE_OUTER);

  auto data_node_find = graph.FindNode("data");
  EXPECT_NE(data_node_find, nullptr);
  graph.ApplySplit(data_node_find, outer_axis.id, inner_axis.id);
}

TEST_F(UtestAscendIR, ApplyMerge) {
  auto ONE = Expression::CreateExpression(1);
  AscGraph graph("test_graph");
  const ExpressionPtr s0 = Expression::CreateExpression("s0");
  Axis &s0_axis = graph.CreateAxis("S0", s0);
  const ExpressionPtr s1 = Expression::CreateExpression("s1");
  Axis &s1_axis = graph.CreateAxis("S1", s1);
  ascir_op::Data data("data", graph);
  auto data_node = ge::NodeUtilsEx::GetNodeFromOperator(data);
  EXPECT_NE(data_node, nullptr);
  data.attr.sched.exec_order = 1;
  data.attr.sched.axis = {s0_axis.id, s1_axis.id};
  data.y.dtype = ge::DT_FLOAT16;
  data.y.format = ge::FORMAT_ND;
  data.y.axis = {s0_axis.id, s1_axis.id};
  data.y.repeats = {s0, s1};
  data.y.strides = {s1, ONE};

  auto merge_axis = graph.MergeAxis({s0_axis.id, s1_axis.id});
  EXPECT_NE(merge_axis, nullptr);
  EXPECT_EQ(merge_axis->type, Axis::AXIS_TYPE_MERGED);

  auto data_node_find = graph.FindNode("data");
  EXPECT_NE(data_node_find, nullptr);
  graph.ApplyMerge(data_node_find, merge_axis->id);
}

TEST_F(UtestAscendIR, ApplyMerge_1) {
  auto ONE = Expression::CreateExpression(1);
  auto ZERO = Expression::CreateExpression(0);
  AscGraph graph("test_graph");
  const auto s0 = Expression::CreateExpression("s0");
  Axis &s0_axis = graph.CreateAxis("S0", s0);
  const auto s1 = Expression::CreateExpression("s1");
  Axis &s1_axis = graph.CreateAxis("S1", s1);
  ascir_op::Data data("data", graph);
  auto data_node = ge::NodeUtilsEx::GetNodeFromOperator(data);
  EXPECT_NE(data_node, nullptr);
  data.attr.sched.exec_order = 1;
  data.attr.sched.axis = {s0_axis.id, s1_axis.id};
  data.y.dtype = ge::DT_FLOAT16;
  data.y.format = ge::FORMAT_ND;
  data.y.axis = {s0_axis.id, s1_axis.id};
  data.y.repeats = {s0, s1};
  data.y.strides = {ZERO, ONE};

  auto merge_axis = graph.MergeAxis({s0_axis.id, s1_axis.id});
  EXPECT_NE(merge_axis, nullptr);
  EXPECT_EQ(merge_axis->type, Axis::AXIS_TYPE_MERGED);

  auto data_node_find = graph.FindNode("data");
  EXPECT_NE(data_node_find, nullptr);
  graph.ApplyMerge(data_node_find, merge_axis->id);
  EXPECT_EQ(data_node_find->attr.sched.axis.size(), 1U);
  EXPECT_EQ(data_node_find->outputs[0].axis.size(), 2U);
}

TEST_F(UtestAscendIR, ApplyReorder) {
  auto ONE = Expression::CreateExpression(1);
  AscGraph graph("test_graph");
  const ExpressionPtr s0 = Expression::CreateExpression("s0");
  Axis &s0_axis = graph.CreateAxis("S0", s0);
  const ExpressionPtr s1 = Expression::CreateExpression("s1");
  Axis &s1_axis = graph.CreateAxis("S1", s1);
  ascir_op::Data data("data", graph);
  auto data_node = ge::NodeUtilsEx::GetNodeFromOperator(data);
  EXPECT_NE(data_node, nullptr);
  data.attr.sched.exec_order = 1;
  data.attr.sched.axis = {s0_axis.id, s1_axis.id};
  data.y.dtype = ge::DT_FLOAT16;
  data.y.format = ge::FORMAT_ND;
  data.y.axis = {s0_axis.id, s1_axis.id};
  data.y.repeats = {s0, s1};
  data.y.strides = {s1, ONE};

  auto data_node_find = graph.FindNode("data");
  EXPECT_NE(data_node_find, nullptr);
  graph.ApplyReorder(data_node_find, {s1_axis.id, s0_axis.id});
}


TEST_F(UtestAscendIR, CreateOptional) {
  AscGraph graph("test_graph");
  (void)graph.CreateOptionalAtt("a", 0U, 1, 2);
  (void)graph.CreateOptionalAtt("a", 0U, 1U, 2U);
  (void)graph.CreateOptionalAtt("a", 0U, 1.0f, 2.0f);
  (void)graph.CreateOptionalAtt("a", 0U, static_cast<int8_t>(1), static_cast<int8_t>(2));
  (void)graph.CreateOptionalAtt("a", 0U, "int32_t", "1", "2");
  auto op_all = graph.GetAllOptional();
  EXPECT_EQ(op_all.size(), 5U);
  EXPECT_EQ(op_all[0]->idx, 0U);
  EXPECT_EQ(op_all[0]->name, "a");
  EXPECT_EQ(op_all[0]->data_type, "int32_t");
  EXPECT_EQ(op_all[0]->min_value, "1");
  EXPECT_EQ(op_all[0]->max_value, "2");
}

TEST_F(UtestAscendIR, Operator_OK) {
  auto ONE = Expression::CreateExpression(1);
  AscGraph graph("test_graph");
  ExpressionPtr s0 = Expression::CreateExpression("s0");
  Axis &s0_axis = graph.CreateAxis("S0", s0);

  ascir_op::Data data("data", graph);
  auto data_node = ge::NodeUtilsEx::GetNodeFromOperator(data);
  EXPECT_NE(data_node, nullptr);
  data.attr.sched.exec_order = 1;
  data.attr.sched.axis = {s0_axis.id};
  data.y.dtype = ge::DT_FLOAT16;
  data.y.format = ge::FORMAT_ND;
  data.y.axis = {s0_axis.id};
  data.y.repeats = {s0};
  data.y.strides = {ONE};

  ascir_op::Abs abs("abs");
  auto abs_node = ge::NodeUtilsEx::GetNodeFromOperator(abs);
  EXPECT_EQ(abs_node, nullptr);
  abs.x = data.y;
  abs_node = ge::NodeUtilsEx::GetNodeFromOperator(abs);
  EXPECT_NE(abs_node, nullptr);

  // find Node
  auto data_node_find = graph.FindNode("data");
  EXPECT_NE(data_node_find, nullptr);
  EXPECT_EQ(data_node_find->attr.sched.exec_order, 1);
  EXPECT_EQ(data_node_find->attr.sched.axis.size(), 1U);
  EXPECT_EQ(data_node_find->attr.sched.axis[0], s0_axis.id);
  EXPECT_EQ(data_node_find->outputs[0].axis.size(), 1);
  EXPECT_EQ(ge::DataType(data_node_find->outputs[0].dtype), ge::DT_FLOAT16);
  EXPECT_EQ(ge::Format(data_node_find->outputs[0].format), ge::FORMAT_ND);
  auto abs_node_find = graph.FindNode("abs");
  EXPECT_NE(abs_node_find, nullptr);

  // GetAllNodes
  int num = 0;
  for (const auto &node : graph.GetAllNodes()) {
    if (num == 0) {
      EXPECT_EQ(node->GetName(), "data");      
      EXPECT_EQ(node->attr.sched.exec_order, 1);
      EXPECT_EQ(node->attr.sched.axis.size(), 1U);
      EXPECT_EQ(node->attr.sched.axis[0], s0_axis.id);
      EXPECT_EQ(node->outputs[0].axis.size(), 1);
      const auto outputs = node->outputs();
      EXPECT_EQ(outputs.size(), 1U);
      EXPECT_NE(outputs[0], nullptr);
      EXPECT_EQ(outputs[0]->axis.size(), 1);
    }
    if (num == 1) {
      EXPECT_EQ(node->inputs.Size(), 1U); 
      EXPECT_EQ(node->inputs[0].axis.size(), 1);
    }
    num++;
  }
  EXPECT_EQ(num, 2);

  // GetAllNodes
  int input_nodes_num = 0;
  for (auto node : graph.GetInputNodes()) {
    if (input_nodes_num == 0) {
      EXPECT_EQ(node->GetName(), "data");      
      EXPECT_EQ(node->attr.sched.exec_order, 1);
      EXPECT_EQ(node->attr.sched.axis.size(), 1U);
      EXPECT_EQ(node->attr.sched.axis[0], s0_axis.id);
      EXPECT_EQ(node->outputs[0].axis.size(), 1);
    }
    input_nodes_num++;
  }
  EXPECT_EQ(input_nodes_num, 1);
  EXPECT_EQ(graph.GetName(), "test_graph");

  // GetAllAxis
  const AscGraph &const_graph = graph;
  const auto all_axis = const_graph.GetAllAxis();
  EXPECT_EQ(all_axis.size(), 1U);
}

TEST_F(UtestAscendIR, Operator_Fail) {
  bool catch_except = false;
  graphStatus status = GRAPH_SUCCESS;
  try {
    ascir_op::Abs abs("abs");
    ascir_op::Output output("output");
    output.x = abs.y;
  } catch (const AscIRException &e) {
    catch_except = true;
    status = e.GetInfo().error_code;
  }
  EXPECT_TRUE(catch_except == true);
  EXPECT_EQ(status, GRAPH_FAILED);
}
